package com.shuigu.ai.factory;

import com.shuigu.ai.assistant.AiCodeGeneratorService;
import com.shuigu.ai.guardrail.PromptSafetyInputGuardrail;
import com.shuigu.ai.guardrail.RetryOutputGuardrail;
import com.shuigu.ai.memory.MongoChatMemoryStore;
import com.shuigu.ai.model.enums.CodeGenTypeEnum;
import com.shuigu.ai.tools.*;
import com.shuigu.common.utils.spring.SpringUtils;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.mcp.client.DefaultMcpClient;
import dev.langchain4j.mcp.client.McpClient;
import dev.langchain4j.mcp.client.transport.McpTransport;
import dev.langchain4j.mcp.client.transport.stdio.StdioMcpTransport;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.tool.ToolProvider;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import okhttp3.internal.http2.ErrorCode;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.swing.text.html.HTML;
import java.io.File;
import java.util.List;

@Slf4j
@Configuration
public class AssistantConfigFactory {

    @Resource
    private MongoChatMemoryStore mongoChatMemoryStore;

    @Resource
    private ToolManager toolManager;

    /**
     * 通过框架的 memoryId 维护会话记忆
     * @return
     */
    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return (memoryId -> MessageWindowChatMemory
                .builder().id(memoryId).maxMessages(20).chatMemoryStore(mongoChatMemoryStore).build());
    }

    /**
     * 不使用框架的 memoryId 维护会话记忆，自己通过 appId 维护会话记忆
     * @param appId
     * @return
     */
    public ChatMemoryProvider getChatMemoryProvider(Long appId) {
        return (memoryId -> MessageWindowChatMemory
                .builder().id(appId).maxMessages(100).chatMemoryStore(mongoChatMemoryStore).build());
    }

    /**
     * 创建新的 AI 服务实例
     */
    public AiCodeGeneratorService createAiCodeGeneratorService(long appId, CodeGenTypeEnum codeGenType) {
        // 根据 appId 构建独立的对话记忆
        ChatMemoryProvider chatMemoryProvider = getChatMemoryProvider(appId);
        ChatMemory chatMemory = chatMemoryProvider.get(appId);
        // 根据代码生成类型选择不同的模型配置
        return switch (codeGenType) {
            // Vue 项目生成使用推理模型
            case VUE_PROJECT -> {
                    // 使用多例模式的 StreamingChatModel 解决并发问题
                    StreamingChatModel customizeStreamingChatModel = SpringUtils.getBean("customizeStreamingChatModel");
                    yield AiServices.builder(AiCodeGeneratorService.class)
                        .streamingChatModel(customizeStreamingChatModel)
                        .chatMemoryProvider(memoryId -> chatMemory)
                        .tools(toolManager.getAllTools())
                        // 最多连续调用 20 次工具,防止AI陷入工具调用的无限循环中
                        .maxSequentialToolsInvocations(20)
                        .inputGuardrails(new PromptSafetyInputGuardrail())
//                        .outputGuardrails(new RetryOutputGuardrail())
                        .hallucinatedToolNameStrategy(toolExecutionRequest -> ToolExecutionResultMessage.from(
                                toolExecutionRequest, "Error: there is no tool called " + toolExecutionRequest.name()
                        ))
                        .build();
            }
            // HTML 和多文件生成使用默认模型
            case HTML, MULTI_FILE -> {
                    // 使用多例模式的 StreamingChatModel 解决并发问题
                    StreamingChatModel openAiStreamingChatModel = SpringUtils.getBean("streamingChatModelPrototype");
                    yield AiServices.builder(AiCodeGeneratorService.class)
                        .streamingChatModel(openAiStreamingChatModel)
                        .chatMemoryProvider(memoryId -> chatMemory)
                        // 输入护轨
                        .inputGuardrails(new PromptSafetyInputGuardrail())
                        // 输出护轨
//                        .outputGuardrails(new RetryOutputGuardrail())
                        .build();
            }
        };
    }

    @Bean
    public ToolProvider mcpToolProvider() {

        McpTransport transport = new StdioMcpTransport.Builder()
                .command(List.of("npx", "-y",
                        "@modelcontextprotocol/server-filesystem",
                        "/Applications/File/myproject/IdeaProject/aiProject/ShuiGu-Langchain4j/shuigu"
                ))
                .logEvents(true)
                .build();

        McpClient mcpClient = new DefaultMcpClient.Builder()
                .transport(transport)
                .build();
        log.info("MCP tool provider initialized:{}", mcpClient.listTools());

        ToolProvider toolProvider = McpToolProvider.builder()
                .mcpClients(List.of(mcpClient))
                .build();

        return toolProvider;
    }
}


